#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <stdbool.h>
#include "utils/utils.h"
#include "matrix_operations.h"
#include "Data_flattening.h"

#define ROWS 28
#define COLS 28
#define MATRIX_SIZE 784
#define HIDDEN_LAYER_SIZE 128
#define OUTPUT_SIZE 10

int main() {
    // 定义矩阵
    double A[ROWS][COLS];
    double B[1][MATRIX_SIZE];
    double D[1][HIDDEN_LAYER_SIZE];
    double F[1][HIDDEN_LAYER_SIZE];
    double G[1][HIDDEN_LAYER_SIZE];
    double H[1][HIDDEN_LAYER_SIZE];
    double L[1][OUTPUT_SIZE];
    double S[1][OUTPUT_SIZE];

    // 用户交互循环
    while (1) {
        char file_path[256];
        printf("Please input filenames:");
        scanf("%s", file_path);

        // 检查文件是否存在
        FILE *fp = fopen(file_path, "rb");
        if (fp == NULL) {
            printf("Invalid file.\n");
            continue;
        }
        fclose(fp);

        // 检查文件格式
        if (!check_pgm_format(file_path)) {
            printf("Invalid file.\n");
            continue;
        }

        // 读取图像文件
        pgm_to_matrix(file_path, ROWS, COLS, A);

        // 将图像矩阵转换为向量
        matrix_to_vector(ROWS, COLS, A, B);

        // 应用第一隐藏层权重和偏置
        Matrix W1 = read_txt_data("input_files/W1.txt");
        double C[ROWS][HIDDEN_LAYER_SIZE];
        matrix_to_vector(W1.rows, W1.cols, (double *)W1.data, (double *)C);
        mul_matrix(1, HIDDEN_LAYER_SIZE, ROWS, B, C, D);
        free_matrix(W1);

        // 添加偏置
        Matrix B1 = read_txt_data("input_files/B1.txt");
        add_matrix(1, HIDDEN_LAYER_SIZE, D, (double *)B1.data, F);
        free_matrix(B1);

        // 激活函数
        active_function(HIDDEN_LAYER_SIZE, F, G);

        // 应用第二隐藏层权重和偏置
        Matrix W2 = read_txt_data("input_files/W2.txt");
        double C2[HIDDEN_LAYER_SIZE][OUTPUT_SIZE];
        matrix_to_vector(W2.rows, W2.cols, (double *)W2.data, (double *)C2);
        mul_matrix(1, OUTPUT_SIZE, HIDDEN_LAYER_SIZE, G, C2, L);
        free_matrix(W2);

        // 添加偏置
        Matrix B2 = read_txt_data("input_files/B2.txt");
        add_matrix(1, OUTPUT_SIZE, L, (double *)B2.data, L);
        free_matrix(B2);

        // SoftMax
        softmax(1, OUTPUT_SIZE, L, S);

        // 输出识别结果
        int digit = 0;
        double max_val = S[0][0];
        for (int i = 1; i < OUTPUT_SIZE; i++) {
            if (S[0][i] > max_val) {
                max_val = S[0][i];
                digit = i;
            }
        }
        printf("%s: %d\n", file_path, digit);

        // 询问是否继续
        char continue_choice;
        printf("Do you want to continue? Please input [Y or N]: ");
        scanf(" %c", &continue_choice);
        if (continue_choice != 'Y' && continue_choice != 'y') {
            break;
        }
    }

    return 0;
}